package com.kingsoft.dc.khaos.module.spark.preprocess.transform

import com.kingsoft.dc.khaos.KhaosContext
import com.kingsoft.dc.khaos.innertype.Schema
import com.kingsoft.dc.khaos.metadata.{Dependency, KhaosStructField}
import com.kingsoft.dc.khaos.module.spark.constants.CustomConstants
import com.kingsoft.dc.khaos.module.spark.metadata.preprocess.transform.BasicInfo
import com.kingsoft.dc.khaos.module.spark.util.DataframeUtils
import com.kingsoft.dc.khaos.util.Logging
import org.apache.spark.sql.{DataFrame, Dataset, Row}
import org.json4s.{DefaultFormats, JObject}
import org.json4s.jackson.JsonMethods.parse

import java.lang.reflect.Method

/**
 * Created by wuxiang on 2022/10/13.
 */
class Custom extends TransformStrategy with Logging {
  private var basicInfo: BasicInfo = _

  override def exec(kc: KhaosContext,
                    module_id: String,
                    config: String,
                    dependences: Seq[Dependency],
                    targets: Seq[Dependency]): Seq[(String, DataFrame)] = {

    // 解析config
    implicit val formats = DefaultFormats
    basicInfo = parse(config, true).extract[BasicInfo]
    var customDataFrame = kc.structData[DataFrame](dependences.head.getEdge())._2
    var customProperties: Map[String, String] = null
    try {
      customProperties = kc.conf.getAllWithPrefix("module.preprocess.transform.custom.").toMap
      log.info("customProperties Properties")
      customProperties.foreach { case (k, v) => log.info(k + "   " + v) }
    } catch {
      case e: Exception =>
        log.error("未读取到custom配置! 改用默认配置")
    }

    val clazz = Class.forName(customProperties.getOrElse(CustomConstants.MODULE_PREPROCESS_TRANSFORM_CUSTOM_MAIN_CALSS, "com.kingsoft.dc.customoperator.RunJar"))
    val execMethod: Method = clazz.getDeclaredMethod(customProperties.getOrElse(CustomConstants.MODULE_PREPROCESS_TRANSFORM_CUSTOM_EXEC_METHOD, "exec"), classOf[Dataset[Row]], classOf[JObject])
    customDataFrame = execMethod.invoke(clazz.newInstance(),  customDataFrame, basicInfo.custom).asInstanceOf[Dataset[Row]]

    // 缓存功能
    if ("true".equals(basicInfo.advanced_options.cache_option.on_off)) {
      customDataFrame = DataframeUtils.cacheDataframe(customDataFrame, this)
    }
    // 分区功能
    if ("true".equals(basicInfo.advanced_options.partition.on_off)) {
      customDataFrame = DataframeUtils.repartionDataframe(customDataFrame, basicInfo.advanced_options.partition.nums.toInt, this)
    }

    addResult(targets.head, customDataFrame)
  }

  override def schema(kc: KhaosContext,
                      module_id: String,
                      config: String,
                      dependences: Seq[Dependency]) = {

    val schema = if (basicInfo.extract_fields.isEmpty) {
      kc.schemaChannel.getSchema(dependences.head.getSource())
    } else {
      val schemaList = basicInfo.extract_fields.map { optInfo =>
        KhaosStructField(optInfo.field, optInfo.data_type)
      }
      new Schema(schemaList)
    }
    schema
  }
}